package com.xiaobao.controller;

import com.xiaobao.service.QianfanService;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

//用于AI的控制器
@RestController
public class ChatController {

    private static final Logger log = LoggerFactory.getLogger(ChatController.class);
    private final QianfanService qianfanService;
    private final ExecutorService executor = Executors.newFixedThreadPool(20);

    public ChatController(QianfanService qianfanService) {
        this.qianfanService = qianfanService;
    }

    @GetMapping(path = "/chat", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter streamChatGet(@RequestParam String message) {
        if (message.isBlank()) {
            return createErrorEmitter("消息不能为空");
        }
        return handleStream(message);
    }

    @PostMapping(path = "/chat", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter streamChatPost(@RequestBody Map<String, String> request) {
        String message = request.getOrDefault("message", "").trim();
        if (message.isBlank()) {
            return createErrorEmitter("消息不能为空");
        }
        return handleStream(message);
    }

    private SseEmitter handleStream(String message) {
        SseEmitter emitter = new SseEmitter(300_000L); // 设置较长的超时时间

        // 标记emitter是否已完成
        final boolean[] isEmitterCompleted = {false};

        // 清理回调
        Runnable cleanupTask = () -> {
            if (!isEmitterCompleted[0]) {
                isEmitterCompleted[0] = true;
                log.info("SSE连接已关闭，message: {}", message);
                try {
                    // 发送结束标记
                    emitter.send(SseEmitter.event()
                            .id("end")
                            .name("message")
                            .data("{\"answer\": \"\"}"));
                } catch (Exception e) {
                    log.warn("发送结束标记失败: {}", e.getMessage());
                }
                emitter.complete();
            }
        };

        // 注册回调
        emitter.onCompletion(cleanupTask);
        emitter.onTimeout(cleanupTask);
        emitter.onError(ex -> {
            log.error("SSE连接异常，message: {}: {}", message, ex.getMessage(), ex);
            try {
                emitter.send(SseEmitter.event()
                        .id("error")
                        .name("message")
                        .data("{\"answer\": \"服务器内部错误，请稍后重试\"}"));
            } catch (IOException e) {
                log.error("发送错误消息失败", e);
            }
            cleanupTask.run();
        });

        executor.submit(() -> {
            try {
                // 发送开始标记
                emitter.send(SseEmitter.event()
                        .id("start")
                        .name("message")
                        .data("{\"answer\": \"\"}"));

                // 调用服务发送数据
                qianfanService.streamAiReply(message, emitter, isEmitterCompleted);

                // 确保emitter已完成
                if (!isEmitterCompleted[0]) {
                    cleanupTask.run();
                }
            } catch (Exception e) {
                log.error("流式回复生成失败，message: {}: {}", message, e.getMessage(), e);
                if (!isEmitterCompleted[0]) {
                    try {
                        emitter.send(SseEmitter.event()
                                .id("error")
                                .name("message")
                                .data("{\"answer\": \"消息处理失败，请重试\"}"));
                    } catch (IOException ex) {
                        log.error("发送错误消息失败", ex);
                    }
                    cleanupTask.run();
                }
            }
        });

        return emitter;
    }

    private SseEmitter createErrorEmitter(String errorMessage) {
        SseEmitter emitter = new SseEmitter();
        try {
            emitter.send(SseEmitter.event()
                    .id("error")
                    .name("message")
                    .data("{\"answer\": \"" + escapeJson(errorMessage) + "\"}"));
            emitter.complete();
        } catch (IOException e) {
            log.error("错误Emitter创建失败: {}", e.getMessage());
        }
        return emitter;
    }

    // 简单的JSON转义方法
    private String escapeJson(String input) {
        if (input == null) {
            return "";
        }
        return input.replace("\\", "\\\\")
                .replace("\"", "\\\"")
                .replace("\n", "\\n")
                .replace("\r", "\\r")
                .replace("\t", "\\t");
    }
}